package tech.xs.common.framework.mq;

import lombok.extern.slf4j.Slf4j;

import java.io.*;
import java.math.BigDecimal;
import java.util.*;

@Slf4j
public class SerialObjectInputStream extends ObjectInputStream {

    private final static Set<String> WHITE_LIST_CLASS = new HashSet<>();

    static {
        WHITE_LIST_CLASS.add(boolean.class.getName());
        WHITE_LIST_CLASS.add(Boolean.class.getName());
        WHITE_LIST_CLASS.add(int.class.getName());
        WHITE_LIST_CLASS.add(Integer.class.getName());
        WHITE_LIST_CLASS.add(long.class.getName());
        WHITE_LIST_CLASS.add(Long.class.getName());
        WHITE_LIST_CLASS.add(float.class.getName());
        WHITE_LIST_CLASS.add(Float.class.getName());
        WHITE_LIST_CLASS.add(double.class.getName());
        WHITE_LIST_CLASS.add(Double.class.getName());
        WHITE_LIST_CLASS.add(Class.class.getName());
        WHITE_LIST_CLASS.add(String.class.getName());
        WHITE_LIST_CLASS.add(Number.class.getName());
        WHITE_LIST_CLASS.add(BigDecimal.class.getName());
        WHITE_LIST_CLASS.add(String.class.getName());
        WHITE_LIST_CLASS.add(Enum.class.getName());
        WHITE_LIST_CLASS.add(ArrayList.class.getName());
        WHITE_LIST_CLASS.add(Hashtable.class.getName());
        WHITE_LIST_CLASS.add(LinkedList.class.getName());
        WHITE_LIST_CLASS.add(HashSet.class.getName());
        WHITE_LIST_CLASS.add(HashMap.class.getName());
        WHITE_LIST_CLASS.add(Hashtable.class.getName());
        WHITE_LIST_CLASS.add(Set.class.getName());
        WHITE_LIST_CLASS.add(Map.class.getName());
        WHITE_LIST_CLASS.add(List.class.getName());
    }

    public SerialObjectInputStream(InputStream in) throws IOException {
        super(in);
    }

    @Override
    protected Class<?> resolveClass(ObjectStreamClass desc) throws IOException, ClassNotFoundException {
        if (WHITE_LIST_CLASS.contains(desc.getName())) {
            return super.resolveClass(desc);
        }
        throw new InvalidClassException("非自定义序列化类,存在安全问题", desc.getName());
    }

    public static void add(String className) {
        log.info("添加Obj序列化白名单 className:" + className);
        WHITE_LIST_CLASS.add(className);
    }

    public static boolean isExist(String className) {
        return WHITE_LIST_CLASS.contains(className);
    }
}
